import pika
import json
import base64
import gzip

from push import push_client
from database import mongodb

class StateTracing:
    def __init__(self, queue_name="trade_state"):
        self.queue_name = queue_name
        self.connection = None
        self.channel = None

        self.push_api = push_client.PushClient()
        self.db_client = mongodb.MongoConnector()

    def init_rabbitmq(self):
        # def callback_rabbitmq(ch, method, properties, body):
        #     print('[x] Received %r' % body)
        self.connection = pika.BlockingConnection(pika.ConnectionParameters('localhost'))
        self.channel = self.connection.channel()
        self.channel.queue_declare(queue=self.queue_name)
        self.channel.queue_purge(queue=self.queue_name)

    def run_tracing(self):
        self.db_client.init_client()
        it = self.channel.consume(queue=self.queue_name, auto_ack=True)
        for _, _, body in it:
            recv_data = json.loads(gzip.decompress(base64.b64decode(body)))
            if recv_data["trade_account"] == "00000000":
                self.db_client.select_db_collection("backtest", "tradestate")
            else:
                self.db_client.select_db_collection("trade", "tradestate")
            # db_collection = self.db_client.get_selected_collection()
            # query_condition = {"code": recv_data["code"], "date": recv_data["date"]}
            result = self.db_client.insert_data([recv_data.copy()])
            if recv_data["trade_account"] == "00000000":
                self.db_client.select_db_collection("backtest", "tradestate_latest")
            else:
                self.db_client.select_db_collection("trade", "tradestate_latest")
            db_collection = self.db_client.get_selected_collection()
            query_condition = {"code": recv_data["code"], "date": recv_data["date"]}
            result = db_collection.update_one(query_condition, {'$set': recv_data.copy()}, upsert=True)
            # print(recv_data)
            # print(result)
        self.db_client.close_client()


if __name__ == "__main__":
    state_tracing = StateTracing()
    state_tracing.init_rabbitmq()
    state_tracing.run_tracing()